import numpy as np
import csv
import os


# 优化后的映射处理器（使用函数属性替代全局变量）
def handleProtocol(inputs):
    if not hasattr(handleProtocol, "mapping"):
        # 初始化协议映射
        handleProtocol.mapping = {}
        handleProtocol.counter = 0

    value = inputs[1]
    if value not in handleProtocol.mapping:
        handleProtocol.mapping[value] = handleProtocol.counter
        handleProtocol.counter += 1
    return handleProtocol.mapping[value]


def handleService(inputs):
    if not hasattr(handleService, "mapping"):
        handleService.mapping = {}
        handleService.counter = 0

    value = inputs[2]
    if value not in handleService.mapping:
        handleService.mapping[value] = handleService.counter
        handleService.counter += 1
    return handleService.mapping[value]


def handleFlag(inputs):
    if not hasattr(handleFlag, "mapping"):
        handleFlag.mapping = {}
        handleFlag.counter = 0

    value = inputs[3]
    if value not in handleFlag.mapping:
        handleFlag.mapping[value] = handleFlag.counter
        handleFlag.counter += 1
    return handleFlag.mapping[value]


def handleLabel(inputs):
    if not hasattr(handleLabel, "mapping"):
        # 初始化已知标签（保持normal.为0）
        handleLabel.mapping = {
            'normal.': 0, 'buffer_overflow.': 1, 'loadmodule.': 2,
            'perl.': 3, 'neptune.': 4, 'smurf.': 5, 'guess_passwd.': 6,
            'pod.': 7, 'teardrop.': 8, 'portsweep.': 9, 'ipsweep.': 10,
            'land.': 11, 'ftp_write.': 12, 'back.': 13, 'imap.': 14,
            'satan.': 15, 'phf.': 16, 'nmap.': 17, 'multihop.': 18,
            'warezmaster.': 19, 'warezclient.': 20, 'spy.': 21, 'rootkit.': 22
        }
        handleLabel.next_index = 23  # 下一个新标签的起始索引

    value = inputs[41]
    if value not in handleLabel.mapping:
        handleLabel.mapping[value] = handleLabel.next_index
        handleLabel.next_index += 1
    return handleLabel.mapping[value]


if __name__ == '__main__':
    input_path = 'data/kddcup.data'
    output_path = "./data/kddcup_data_numerization.txt"

    # 检查输入文件是否存在
    if not os.path.exists(input_path):
        raise FileNotFoundError(f"数据集文件不存在: {input_path}")

    with open(input_path, 'r') as data_original, \
            open(output_path, 'w', newline='') as data_numerization:

        csv_reader = csv.reader(data_original)
        csv_writer = csv.writer(data_numerization)

        for row in csv_reader:
            # 跳过空行
            if not row:
                continue

            # 复制新行用于修改
            new_row = list(row)

            # 处理非数值特征
            try:
                new_row[1] = handleProtocol(row)
                new_row[2] = handleService(row)
                new_row[3] = handleFlag(row)
                new_row[41] = handleLabel(row)
            except IndexError as e:
                print(f"行格式错误: {row[:5]}... 跳过")
                continue

            # 转换数值类型
            for i in range(len(new_row)):
                if i in (1, 2, 3, 41):  # 已处理的分类列
                    continue
                try:
                    new_row[i] = float(new_row[i])
                except ValueError:
                    pass  # 保持非数值数据不变

            csv_writer.writerow(new_row)

    print('数值化完成!')
    print(f'协议类型数量: {handleProtocol.counter}')
    print(f'服务类型数量: {handleService.counter}')
    print(f'标志类型数量: {handleFlag.counter}')
    print(f'标签类型数量: {handleLabel.next_index}')